import numpy as np


def generate_9x9_orthogonal_latin_squares():
    """
    生成 9×9 的正交拉丁方
    """
    # 第一个拉丁方
    latin_square1 = np.zeros((9, 9), dtype=int)
    for i in range(9):
        for j in range(9):
            latin_square1[i, j] = (i + j) % 9

    # 第二个拉丁方
    latin_square2 = np.zeros((9, 9), dtype=int)
    for i in range(9):
        for j in range(9):
            latin_square2[i, j] = (i * 2 + j) % 9

    return latin_square1, latin_square2


def permute_data(data, latin_square1, latin_square2):
    """
    根据正交拉丁方对一维数据进行置乱
    :param data: 一维数据
    :param latin_square1: 第一个拉丁方
    :param latin_square2: 第二个拉丁方
    :return: 置乱后的数据
    """
    n = int(np.sqrt(len(data)))
    data_matrix = data.reshape((n, n))
    permuted_matrix = np.zeros((n, n), dtype=data.dtype)
    for i in range(n):
        for j in range(n):
            new_i = latin_square1[i, j]
            new_j = latin_square2[i, j]
            permuted_matrix[new_i, new_j] = data_matrix[i, j]
    return permuted_matrix.flatten()


def recover_data(permuted_data, latin_square1, latin_square2):
    """
    根据正交拉丁方对置乱后的数据进行恢复
    :param permuted_data: 置乱后的数据
    :param latin_square1: 第一个拉丁方
    :param latin_square2: 第二个拉丁方
    :return: 恢复后的数据
    """
    n = int(np.sqrt(len(permuted_data)))
    permuted_matrix = permuted_data.reshape((n, n))
    recovered_matrix = np.zeros((n, n), dtype=permuted_data.dtype)
    for i in range(n):
        for j in range(n):
            new_i = latin_square1[i, j]
            new_j = latin_square2[i, j]
            recovered_matrix[i, j] = permuted_matrix[new_i, new_j]
    return recovered_matrix.flatten()


if __name__ == "__main__":
    # 示例一维数据
    data = np.arange(1, 82)
    print("原始数据:", data)

    # 生成 9×9 正交拉丁方
    latin_square1, latin_square2 = generate_9x9_orthogonal_latin_squares()
    print("生成的第一个 9×9 拉丁方:")
    print(latin_square1)
    print("生成的第二个 9×9 拉丁方:")
    print(latin_square2)

    # 置乱数据
    permuted_data = permute_data(data, latin_square1, latin_square2)
    print("置乱后的数据:", permuted_data)

    # 恢复数据
    recovered_data = recover_data(permuted_data, latin_square1, latin_square2)
    print("恢复后的数据:", recovered_data)
